#! /usr/bin/env python
#coding=utf-8

import concurrent.futures

from macho_file_base import MachOFileBase
from macho_walker import MachOWalker

class MachOFileDeps(MachOFileBase):
    def __init__(self, file, prefix, mgr):
        super(MachOFileDeps, self).__init__(file, prefix)
        self["deps"] = []
        self["dependedBy"] = []

        self["deps_indirect"] = []
        self["dependedBy_indirect"] = []
        self["deps_total"] = 0
        self["dependedBy_total"] = 0

        self["mgr"] = mgr

    def __eq__(self, other):
        if not isinstance(other, MachOFileDeps):
            return NotImplemented

        return self["id"] == other["id"]#and self["name"] == other["name"]

    def dependsOn(self, mod):
        for dep in self["deps"]:
            if dep["callee"] == mod:
                return True
        return False

    def getAllDependedModules(self):
        res = []
        for dep in self["deps"]:
            res.append(dep["callee"])
        return res + self["deps_indirect"]

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "%s:%d deps(%d) depsTotal(%d) dependedBy(%d)" % (self["name"], self["id"], len(self["deps"]), len(self["deps"]) + len(self["deps_indirect"]), len(self["dependedBy"]))

class DependencyBase(dict):
    def __init__(self, caller, callee):
        self["caller_id"] = caller["id"]
        self["callee_id"] = callee["id"]
        self["caller"] = caller
        self["callee"] = callee
        self["calls"] = 0

    def __eq__(self, other):
        if not isinstance(other, DependencyBase):
            return NotImplemented

        return self["caller_id"] == other["caller_id"] and self["callee_id"] == other["callee_id"]

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "%d:(%s[%d] -%d-> %s[%d])" % (self["id"], self["caller"]["name"], self["caller"]["id"], self["calls"], self["callee"]["name"], self["callee"]["id"])

class MachOFileMgr(object):
    def __init__(self, dylib_shared_cache=None, fileClass=None, dependenceClass = None):
        self._files = []
        self._path_dict = {}
        if fileClass:
            self.__dylibFileClass = fileClass
        else:
            self.__dylibFileClass = MachOFileDeps

        if dependenceClass:
            self._dependenceClass = dependenceClass
        else:
            self._dependenceClass = DependencyBase
        self._fileIdx = 1

        if dylib_shared_cache:
            walker = MachOWalker(dylib_shared_cache)
        else:
            walker = MachOWalker()
        self._prefix = walker.get_path()

    def get_root_path(self):
        return self._prefix

    def scan_all_files(self):
        walker = MachOWalker(self._prefix)

        self._scan_all_dylib_files(walker)

        self.__build_deps_tree()

    def __build_deps_tree(self):
        self._build_deps_tree()

        self._maxDepth = 0
        self._maxTotalDepends = 0

        print("Build indirect dependence tree for %d dylib files now ..." % len(self._files))

        for mod in self._files:
            mod["_recursiveFinished"] = False
        #for mod in self._dylibFiles:
        #	self.__update_indirect_deps_recursive(mod)
        for mod in self._files:
            mod["_recursiveFinished"] = False
        #for mod in self._dylibFiles:
        #	self.__update_indirect_dependedBy_recursive(mod)
        for mod in self._files:
            del mod["_recursiveFinished"]

    def add_one_file(self, file):
        # Append to array in order
        file["id"] = self._fileIdx
        self._fileIdx = self._fileIdx + 1
        self._files.append(file)

        # Add to dictionary with path as key
        self._path_dict[file["name"]] = file

    def _scan_all_dylib_files(self, walker):
        print("Scanning %d Dylib files now ..." % len(walker.get_all_files()))
        for f in walker.get_all_files():
            dylib = self.__dylibFileClass(f, self._prefix, self)
            if dylib["name"] in self._path_dict:
                print("Warning: duplicate " + dylib.get_file() + ' skipped.')
                continue

            self.add_one_file(dylib)

    def _build_deps_tree(self):
        print("Build dependence tree for %d dylib files now ..." % len(self._files))
        self.__parse_depends()
        cnt = 0
        for file in self._files:
            cnt += self.__setup_depends(file)
        print("    Got %d dependencies" % cnt)

    def __setup_depends(self, mod):
        startingId = mod["id"] << 16
        cnt = 0
        for dep in mod["deps"]:
            cnt = cnt + 1
            dep["id"] = startingId + cnt
            dep["callee"]["dependedBy"].append(dep)
        return cnt

    def __parse_depends_task(mod):
        mgr = mod["mgr"]
        for lib in mod.library_depends():
            dep_file = mgr.get_file_by_path(lib)
            if not dep_file:
                print("Warning: can not find depended library [" + lib + "] for " + mod["name"])
                continue
            dep = mgr._dependenceClass(mod, dep_file)
            mod["deps"].append(dep)

    def __parse_depends(self):
        with concurrent.futures.ThreadPoolExecutor() as executor:
            executor.map(MachOFileMgr.__parse_depends_task, self._files, chunksize=10)

    def get_file_by_path(self, path):
        if path in self._path_dict:
            return self._path_dict[path]
        return None

    def get_file_by_idx(self, idx):
        if idx < 1 or idx > len(self._files):
            return None
        return self._files[idx - 1]

    def get_all(self):
        return self._files

    def __update_indirect_dependedBy_recursive(self, mod):
        # Already finished
        if mod["_recursiveFinished"]:
            return mod["dependedBy_depth"]

        maxDepth = 0
        for item in mod["dependedBy"]:
            # update caller first
            caller = item["caller"]
            depth = self.__update_indirect_dependedBy_recursive(caller)
            if depth > maxDepth:
                maxDepth = depth
            for dep in caller["dependedBy"]:
                grand_caller = dep["caller"]
                if grand_caller.dependsOn(mod):
                    continue
                if grand_caller in mod["dependedBy_indirect"]:
                    continue
                mod["dependedBy_indirect"].append(grand_caller)
            for dep in caller["dependedBy_indirect"]:
                if dep.dependsOn(mod):
                    continue
                if dep in mod["dependedBy_indirect"]:
                    continue
                mod["dependedBy_indirect"].append(dep)

        if len(mod["dependedBy"]) > 0:
            maxDepth = maxDepth + 1

        mod["_recursiveFinished"] = True
        mod["dependedBy_depth"] = maxDepth

        if maxDepth > self._maxDepth:
            self._maxDepth = maxDepth
        depsTotal = len(mod["dependedBy"]) + len(mod["dependedBy_indirect"])
        if depsTotal > self._maxTotalDepends:
            self._maxTotalDepends = depsTotal

        mod["dependedBy_total"] = depsTotal

        return maxDepth

    def __update_indirect_deps_recursive(self, mod):
        # Already finished
        if mod["_recursiveFinished"]:
            return mod["depth"]

        maxDepth = 0
        for item in mod["deps"]:
            # update child first
            child = item["callee"]
            depth = self.__update_indirect_deps_recursive(child)
            if depth > maxDepth:
                maxDepth = depth
            for dep in child["deps"]:
                if mod.dependsOn(dep["callee"]):
                    continue
                if dep["callee"] in mod["deps_indirect"]:
                    continue
                mod["deps_indirect"].append(dep["callee"])
            for dep in child["deps_indirect"]:
                if mod.dependsOn(dep):
                    continue
                if dep in mod["deps_indirect"]:
                    continue
                mod["deps_indirect"].append(dep)

        if len(mod["deps"]) > 0:
            maxDepth = maxDepth + 1

        mod["_recursiveFinished"] = True
        mod["depth"] = maxDepth

        if maxDepth > self._maxDepth:
            self._maxDepth = maxDepth
        depsTotal = len(mod["deps"]) + len(mod["deps_indirect"])
        if depsTotal > self._maxTotalDepends:
            self._maxTotalDepends = depsTotal

        mod["deps_total"] = depsTotal

        return maxDepth

if __name__ == '__main__':
    mgr = MachOFileMgr()
    mgr.scan_all_files()
    f = mgr.get_file_by_path("/usr/lib/libutil.dylib")
    print("Get libutil now ...")

    print(f)
    #print(mgr.get_all())
    print(len(f["dependedBy"]))
    for idx, f in enumerate(mgr.get_all()):
        if idx + 1 != f["id"]:
            print(f["name"])
